# Spectral GCN + Attention Recovery + LSTM
# This code trains and tests the GNN model for the COVID-19 infection prediction in Tokyo
# Author: Jiawei Xue, August 26, 2021
# Step 1: read and pack the traning and testing data
# Step 2: training epoch, training process, testing
# Step 3: build the model = spectral GCN + Attention Recovery + LSTM
# Step 4: main function
# Step 5: evaluation
# Step 6: visualization

import os
import csv
import json
import copy
import time
import random
import string
import argparse
import numpy as np
import pandas as pd
import geopandas as gpd
import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
import matplotlib.pyplot as plt
from matplotlib import pyplot as plt
import torch.nn.functional as F
from spectral_T3_GCN_memory_light import SpecGCN
from spectral_T3_GCN_memory_light import SpecGCN_LSTM
#torch.set_printoptions(precision=8)
#hyperparameter for the setting
X_day, Y_day = 21,21
#START_DATE, END_DATE = '20200414','20210207'
#START_DATE, END_DATE = '20200808','20210603'
START_DATE, END_DATE = '20200720','20210515'
WINDOW_SIZE = 7

#hyperparameter for the learning
DROPOUT, ALPHA = 0.50, 0.20
NUM_EPOCHS, BATCH_SIZE, LEARNING_RATE = 100, 8, 0.0001
HIDDEN_DIM_1, OUT_DIM_1, HIDDEN_DIM_2 = 6,4,2
infection_normalize_ratio = 100.0
web_search_normalize_ratio = 100.0
train_ratio = 0.7
validate_ratio = 0.1

Step 1: read and pack the training and testing data

#1.total period (mobility+text): 
#from 20200201 to 20210620: (29+31+30+31+30+31+31+30+31+30+31)+(31+28+31+30+31+20)\
#= 335 + 171 = 506;
#2.number of zones: 23;
#3.infection period:
#20200331 to 20210620: (1+30+31+30+31+31+30+31+30+31)+(31+28+31+30+31+20) = 276 + 171 = 447.

#1. Mobility: functions 1.2 to 1.7
#2. Text:  functions 1.8 to 1.14
#3. InfectionL:  functions 1.15 
#4. Preprocess:  functions 1.16 to 1.24
#5. Learn:  functions 1.25 to 1.26

#function 1.1
#get the central areas of Tokyo (e.g., the Special wards of Tokyo)
#return: a 23 zone shapefile
def read_tokyo_23():
    folder = "/data/HSEES/xue/xue_codes/disease_prediction_ml/gml_code/present_model_version10/tokyo_23" 
    file = "tokyo_23zones.shp"
    path = os.path.join(folder,file) 
    data = gpd.read_file(path)   
    return data

##################1.Mobility#####################
#function 1.2
#get the average of two days' mobility (infection) records
def mob_inf_average(data, key1, key2):
    new_record = dict()
    record1, record2 = data[key1], data[key2]
    for i in record1:
        if i in record2:
            new_record[i] = (record1[i]+record2[i])/2.0
    return new_record

#function 1.3
#get the average of multiple days' mobility (infection) records
def mob_inf_average_multiple(data, keyList):
    new_record = dict()
    num_day = len(keyList)
    for i in range(num_day):
        record = data[keyList[i]]
        for zone_id in record:
            if zone_id not in list(new_record.keys()):
                new_record[zone_id] = record[zone_id]
            else:
                new_record[zone_id] += record[zone_id]
    for new_record_key in new_record:
        new_record[new_record_key] = new_record[new_record_key]*1.0/num_day
    return new_record

#function 1.4
#generate the dateList: [20200101, 20200102, ..., 20211231]
def generate_dateList():
    yearList = ["2020","2021"]
    monthList = ["0"+str(i+1) for i in range(9)] + ["10","11","12"]
    dayList = ["0"+str(i+1) for i in range(9)] + [str(i) for i in range(10,32)]
    day_2020_num = [31,29,31,30,31,30,31,31,30,31,30,31]
    day_2021_num = [31,28,31,30,31,30,31,31,30,31,30,31]
    date_2020, date_2021 = list(), list()
    for i in range(12):
        for j in range(day_2020_num[i]):
              date_2020.append(yearList[0] + monthList[i] + dayList[j])
        for j in range(day_2021_num[i]):
              date_2021.append(yearList[1] + monthList[i] + dayList[j])
    date_2020_2021 = date_2020 + date_2021
    return date_2020_2021

#function 1.5
#smooth the mobility (infection) data using the neighborhood average
#under a given window size 
#dateList: [20200101, 20200102, ..., 20211231]
def mob_inf_smooth(data, window_size, dateList):
    data_copy = copy.copy(data)
    data_key_list = list(data_copy.keys())
    for data_key in data_key_list:
        left = int(max(dateList.index(data_key)-(window_size-1)/2, 0))
        right = int(min(dateList.index(data_key)+(window_size-1)/2, len(dateList)-1))
        potential_neighbor = dateList[left:right+1]
        neighbor_data_key = list(set(data_key_list).intersection(set(potential_neighbor)))
        data_average = mob_inf_average_multiple(data_copy, neighbor_data_key)
        data[data_key] =  data_average
    return data

#function 1.6
#set the mobility (infection) of one day as zero
def mob_inf_average_null(data, key1, key2):
    new_record = dict()
    record1, record2 = data[key1], data[key2]
    for i in record1:
        if i in record2:
            new_record[i] = 0
    return new_record

#function 1.7
#read the mobility data from "mobility_feature_20200201.json"...
#return: all_mobility:{"20200201":{('123','123'):12345,...},...}
#20200201 to 20210620: 506 days
def read_mobility_data(jcode23):
    all_mobility = dict()
    mobilityFilePath = "/data/HSEES/xue/xue_codes/disease_prediction_ml/gml_code/"+\
    "present_model_version10/mobility_20210804"
    mobilityNameList = os.listdir(mobilityFilePath)
    for i in range(len(mobilityNameList)):
        day_mobility = dict()
        file_name = mobilityNameList[i] 
        if "20" in file_name:
            day = (file_name.split("_")[2]).split(".")[0]  #get the day
            file_path = mobilityFilePath + '/' + file_name
            f = open(file_path,)
            df_file = json.load(f)   #read the mobility file
            f.close()
            for key in df_file:
                origin, dest = key.split("_")[0], key.split("_")[1]
                if origin in jcode23 and dest in jcode23:
                    if origin == dest:
                        day_mobility[(origin, dest)] = 0.0     #ignore the inner-zone flow
                    else:
                        day_mobility[(origin, dest)] = df_file[key] 
            all_mobility[day] = day_mobility
    #missing data
    all_mobility["20201128"] = mob_inf_average(all_mobility,"20201127","20201129")
    all_mobility["20210104"] = mob_inf_average(all_mobility, "20210103","20210105")
    return all_mobility

##################2.Text#####################
#function 1.8
#get the average of two days' infection records
def text_average(data, key1, key2):
    new_record = dict()
    record1, record2 = data[key1], data[key2]
    for i in record1:
        if i in record2:
            zone_record1, zone_record2 = record1[i], record2[i]
            new_zone_record = dict()
            for j in zone_record1:
                if j in zone_record2:
                    new_zone_record[j] = (zone_record1[j] + zone_record2[j])/2.0
            new_record[i] = new_zone_record
    return new_record

#function 1.9
#get the average of multiple days' text records
def text_average_multiple(data, keyList):
    new_record = dict()
    num_day = len(keyList)
    for i in range(num_day):
        record = data[keyList[i]]
        for zone_id in record:                           #zone_id
            if zone_id not in new_record:
                new_record[zone_id] = dict()   
            for j in record[zone_id]:                    #symptom
                if j not in new_record[zone_id]:
                    new_record[zone_id][j] = record[zone_id][j]
                else: 
                    new_record[zone_id][j] += record[zone_id][j]
    for zone_id in new_record:
        for j in new_record[zone_id]:
            new_record[zone_id][j] = new_record[zone_id][j]*1.0/num_day 
    return new_record

#function 1.10
#smooth the text data using the neighborhood average
#under a given window size 
def text_smooth(data, window_size, dateList):
    data_copy = copy.copy(data)
    data_key_list = list(data_copy.keys())
    for data_key in data_key_list:
        left = int(max(dateList.index(data_key)-(window_size-1)/2, 0))
        right = int(min(dateList.index(data_key)+(window_size-1)/2, len(dateList)-1))
        potential_neighbor = dateList[left:right+1]
        neighbor_data_key = list(set(data_key_list).intersection(set(potential_neighbor)))
        data_average = text_average_multiple(data_copy, neighbor_data_key)
        data[data_key] =  data_average
    return data

#function 1.11
#read the number of user points
def read_point_json():
    with open('user_point/mobility_user_point.json') as point1:
        user_point1 = json.load(point1)
    with open('user_point/mobility_user_point_20210812.json') as point2:
        user_point2 = json.load(point2)
    user_point_all = dict()
    for i in user_point1:
        user_point_all[i] = user_point1[i]
    for i in user_point2:
        user_point_all[i] = user_point2[i]
    user_point_all["20201128"] = user_point_all["20201127"]  #data missing
    user_point_all["20210104"] = user_point_all["20210103"]  #data missing
    return user_point_all

#function 1.12
#normalize the text search by the number of user points.
def normalize_text_user(all_text, user_point_all):
    for day in all_text:
        if day in user_point_all:
            num_user = user_point_all[day]["num_user"]
            all_text_day_new = dict()
            all_text_day = all_text[day]
            for zone in all_text_day:
                if zone not in all_text_day_new:
                    all_text_day_new[zone] = dict()
                for sym in all_text_day[zone]:
                    all_text_day_new[zone][sym] = all_text_day[zone][sym]*1.0/num_user
            all_text[day] = all_text_day_new
    return all_text
    
#function 1.13
#read the text data
#20200201 to 20210620: 506 days
#all_text = {"20200211":{"123":{"code":3,"fever":2,...},...},...}
def read_text_data(jcode23):
    all_text = dict()
    textFilePath = "/data/HSEES/xue/xue_codes/disease_prediction_ml/gml_code/"+\
    "present_model_version10/text_20210804"
    textNameList = os.listdir(textFilePath)
    for i in range(len(textNameList)):
        day_text = dict()
        file_name = textNameList[i]
        if "20" in file_name:
            day = (file_name.split("_")[2]).split(".")[0]
            file_path = textFilePath + "/" + file_name
            f = open(file_path,)
            df_file = json.load(f)   #read the mobility file
            f.close()
            new_dict = dict()
            for key in df_file:
                if key in jcode23:
                    new_dict[key] = {key1:df_file[key][key1]*1.0*web_search_normalize_ratio for key1 in df_file[key]}
                    #new_dict[key] = df_file[key]*WEB_SEARCH_RATIO
            all_text[day] = new_dict
    all_text["20201030"] = text_average(all_text, "20201029", "20201031") #data missing
    return all_text

#function 1.14
#perform the min-max normalization for the text data.
def min_max_text_data(all_text,jcode23):
    #calculate the min_max
    #region_key: sym: [min,max]
    text_list = list(['痛み', '頭痛', '咳', '下痢', 'ストレス', '不安', \
                    '腹痛', 'めまい', '吐き気', '嘔吐', '筋肉痛', '動悸', \
                    '副鼻腔炎', '発疹', 'くしゃみ', '倦怠感', '寒気', '脱水', \
                    '中咽頭', '関節痛', '不眠症', '睡眠障害', '鼻漏', '片頭痛', \
                    '多汗症', 'ほてり', '胸痛', '発汗', '無気力', '呼吸困難', \
                    '喘鳴', '目の痛み', '体の痛み', '無嗅覚症', '耳の痛み', \
                    '錯乱', '見当識障害', '胸の圧迫感', '鼻の乾燥', '耳感染症', \
                    '味覚消失', '上気道感染症', '眼感染症', '食欲減少'])
    region_sym_min_max = dict()
    for key in jcode23:                          #initialize
        region_sym_min_max[key] = dict()
        for sym in text_list:
            region_sym_min_max[key][sym] = [1000000,0]  #min, max
    for day in all_text:                            #update
        for key in jcode23:
            for sym in text_list:
                if sym in all_text[day][key]:
                    count = all_text[day][key][sym]
                    if count < region_sym_min_max[key][sym][0]:
                        region_sym_min_max[key][sym][0] = count
                    if count > region_sym_min_max[key][sym][1]:
                        region_sym_min_max[key][sym][1] = count
    #print ("region_sym_min_max",region_sym_min_max)
    for key in jcode23:              #normalize
        for sym in text_list:
            min_count,max_count=region_sym_min_max[key][sym][0],region_sym_min_max[key][sym][1]
            for day in all_text:
                if sym in all_text[day][key]:
                    if max_count-min_count == 0:
                        all_text[day][key][sym] = 1
                    else:
                        all_text[day][key][sym] = (all_text[day][key][sym]-min_count)*1.0/(max_count-min_count)
                        #print("all_text[day][key][sym]",all_text[day][key][sym])
    return all_text

##################3.Infection#####################
#function 1.15
#read the infection data
#20200331 to 20210620: (1+30+31+30+31+31+30+31+30+31)+(31+28+31+30+31+20) = 276 + 171 = 447.
#all_infection = {"20200201":{"123":1,"123":2}}
def read_infection_data(jcode23):
    all_infection = dict()
    infection_path = "/data/HSEES/xue/xue_codes/disease_prediction_ml/gml_code/"+\
    "present_model_version10/patient_20210725.json"
    f = open(infection_path,)
    df_file = json.load(f)   #read the mobility file
    f.close()
    for zone_id in df_file:
        for one_day in df_file[zone_id]:
            daySplit = one_day.split("/") 
            year, month, day = daySplit[0], daySplit[1], daySplit[2]
            if len(month) == 1:
                month = "0" + month
            if len(day) == 1:
                day = "0" + day
            new_date = year + month + day
            if str(zone_id[0:5]) in jcode23:
                if new_date not in all_infection:
                    all_infection[new_date] = {zone_id[0:5]:df_file[zone_id][one_day]*1.0/infection_normalize_ratio}
                else:
                    all_infection[new_date][zone_id[0:5]] = df_file[zone_id][one_day]*1.0/infection_normalize_ratio
    #missing
    date_list = [str(20200316+i) for i in range(15)]
    for date in date_list:
        all_infection[date] = mob_inf_average(all_infection,'20200401','20200401')
    all_infection['20200514'] = mob_inf_average(all_infection,'20200513','20200515')
    all_infection['20200519'] = mob_inf_average(all_infection,'20200518','20200520')
    all_infection['20200523'] = mob_inf_average(all_infection,'20200522','20200524')
    all_infection['20200530'] = mob_inf_average(all_infection,'20200529','20200601')
    all_infection['20200531'] = mob_inf_average(all_infection,'20200529','20200601')
    all_infection['20201231'] = mob_inf_average(all_infection,'20201230','20210101')
    all_infection['20210611'] = mob_inf_average(all_infection,'20210610','20210612')
    #outlier
    all_infection['20200331'] = mob_inf_average(all_infection,'20200401','20200401')
    all_infection['20200910'] = mob_inf_average(all_infection,'20200909','20200912')
    all_infection['20200911'] = mob_inf_average(all_infection,'20200909','20200912')
    all_infection['20200511'] = mob_inf_average(all_infection,'20200510','20200512')
    all_infection['20201208'] = mob_inf_average(all_infection,'20201207','20201209')
    all_infection['20210208'] = mob_inf_average(all_infection,'20210207','20210209')
    all_infection['20210214'] = mob_inf_average(all_infection,'20210213','20210215')
    #calculate the subtraction
    all_infection_subtraction = dict()
    all_infection_subtraction['20200331'] = all_infection['20200331']
    all_keys = list(all_infection.keys())
    all_keys.sort()
    for i in range(len(all_keys)-1):
        record = dict()
        for j in all_infection[all_keys[i+1]]:
            record[j] = all_infection[all_keys[i+1]][j] - all_infection[all_keys[i]][j]
        all_infection_subtraction[all_keys[i+1]] = record
    return all_infection_subtraction, all_infection

##################4.Preprocess#####################
#function 1.16
#ensemble the mobility, text, and infection.
#all_mobility = {"20200201":{('123','123'):12345,...},...}
#all_text = {"20200201":{"123":{"cold":3,"fever":2,...},...},...}
#all_infection = {"20200316":{"123":1,"123":2}}
#all_x_y = {"0":[[mobility_1,text_1, ..., mobility_x_day,text_x_day], [infection_1,...,infection_y_day],\
#[infection_1,...,infection_x_day]],0}
#x_days, y_days: use x_days to predict y_days
def ensemble(all_mobility, all_text, all_infection, x_days, y_days, all_day_list):
    all_x_y = dict()
    for j in range(len(all_day_list) - x_days - y_days + 1):
        x_sample, y_sample, x_sample_infection = list(), list(), list()                   
        #add the data from all_day_list[0+j] to all_day_list[x_days-1+j]
        for k in range(x_days):
            day = all_day_list[k + j]
            x_sample.append(all_mobility[day])
            x_sample.append(all_text[day])  
            x_sample_infection.append(all_infection[day])             #concatenate with the infection data                       
        #add the data from all_day_list[x_days+j] to all_day_list[x_days+y_day-1+j]
        for k in range(y_days):
            day = all_day_list[x_days + k + j]
            y_sample.append(all_infection[day]) 
        all_x_y[str(j)] = [x_sample, y_sample, x_sample_infection,j]                          
    return all_x_y

#function 1.17
#split the data by train/validate/test = train_ratio/validation_ratio/(1-train_ratio-validation_ratio)
def split_data(all_x_y, train_ratio, validation_ratio):
    all_x_y_key = list(all_x_y.keys())
    n = len(all_x_y_key)
    n_train, n_validate = round(n*train_ratio), round(n*validation_ratio)
    n_test = n-n_train-n_validate
    train_key = [all_x_y[str(i)] for i in range(n_train)]
    validate_key = [all_x_y[str(i+n_train)] for i in range(n_validate)]
    test_key = [all_x_y[str(i+n_train+n_validate)] for i in range(n_test)]
    return train_key, validate_key, test_key

##function 1.18
#the second data split method
#split the data by train/validate/test = train_ratio/validation_ratio/(1-train_ratio-validation_ratio)
def split_data_2(all_x_y, train_ratio, validation_ratio):
    all_x_y_key = list(all_x_y.keys())
    n = len(all_x_y_key)
    n_train, n_validate = round(n*train_ratio), round(n*validation_ratio)
    n_test = n-n_train-n_validate
    train_list, validate_list = list(), list()
    
    train_validate_key = [all_x_y[str(i)] for i in range(n_train+n_validate)]
    train_key, validate_key = list(), list()
    for i in range(len(train_validate_key)):
        if i % 9 == 8:
            validate_key.append(all_x_y[str(i)])
            validate_list.append(i)
        else:
            train_key.append(all_x_y[str(i)])
            train_list.append(i)
    test_key = [all_x_y[str(i+n_train+n_validate)] for i in range(n_test)]
    return train_key, validate_key, test_key, train_list, validate_list

##function 1.19
#the third data split method
#split the data by train/validate/test = train_ratio/validation_ratio/(1-train_ratio-validation_ratio)
def split_data_3(all_x_y, train_ratio, validation_ratio):
    all_x_y_key = list(all_x_y.keys())
    n = len(all_x_y_key)
    n_train, n_validate = round(n*train_ratio), round(n*validation_ratio)
    n_test = n - n_train  - n_validate
    train_list, validate_list = list(), list()
    
    train_validate_key = [all_x_y[str(i)] for i in range(n_train + n_validate)]
    train_key, validate_key = list(), list()
    for i in range(len(train_validate_key)):
        if (n_train + n_validate-i) % 2 == 0 and (n_train + n_validate-i) <= 2*n_validate:
            validate_key.append(all_x_y[str(i)])
            validate_list.append(i)
        else:
            train_key.append(all_x_y[str(i)])
            train_list.append(i)
    test_key = [all_x_y[str(i+n_train+n_validate)] for i in range(n_test)]
    return train_key, validate_key, test_key, train_list, validate_list

##function 1.20
#find the mobility data starting from the day, which is x_days before the start_date
#start_date = "20200331", x_days = 7
def sort_date(all_mobility, start_date, x_days): 
    mobility_date_list = list(all_mobility.keys())
    mobility_date_list.sort()
    idx = mobility_date_list.index(start_date)
    mobility_date_cut = mobility_date_list[idx-x_days:] 
    return mobility_date_cut

#function 1.21
#find the mobility data starting from the day, which is x_days before the start_date,
#ending at the day, which is y_days after the end_date
#start_date = "20200331", x_days = 7
def sort_date_2(all_mobility, start_date, x_days, end_date, y_days): 
    mobility_date_list = list(all_mobility.keys())
    mobility_date_list.sort()
    idx = mobility_date_list.index(start_date)
    idx2 = mobility_date_list.index(end_date)
    mobility_date_cut = mobility_date_list[idx-x_days:idx2+y_days] 
    return mobility_date_cut

#function 1.22
#get the mappings from zone id to id, text id to id.
#get zone_text_to_idx 
def get_zone_text_to_idx(all_infection):
    zone_list = list(set(all_infection["20200401"].keys()))
    text_list = list(['痛み', '頭痛', '咳', '下痢', 'ストレス', '不安', \
                    '腹痛', 'めまい'])
    zone_list.sort()
    zone_dict = {str(zone_list[i]):i for i in range(len(zone_list))}
    text_dict = {str(text_list[i]):i for i in range(len(text_list))}
    return zone_dict, text_dict

#function 1.23
#change the data format to matrix
#zoneid_to_idx = {"13101":0, "13102":1, ..., "13102":22}
#sym_to_idx = {"cough":0}
#mobility: {('13101', '13101'): 709973, ...}
#text: {'13101': {'痛み': 51,...},...}  text
#infection: {'13101': 50, '13102': 137, '13103': 401,...} 
#data_type = {"mobility", "text", "infection"}
def to_matrix(zoneid_to_idx, sym_to_idx, input_data, data_type):
    n_zone, n_text = len(zoneid_to_idx), len(sym_to_idx)
    if data_type == "mobility":
        result = np.zeros((n_zone, n_zone))
        for key in input_data:    
            from_id, to_id = key[0], key[1]
            from_idx, to_idx = zoneid_to_idx[from_id], zoneid_to_idx[to_id]
            result[from_idx][to_idx] += input_data[key]
    if data_type == "text":
        result = np.zeros((n_zone, n_text))
        for key1 in input_data:
            for key2 in input_data[key1]:
                if key1 in list(zoneid_to_idx.keys()) and key2 in list(sym_to_idx.keys()):
                    zone_idx, text_idx = zoneid_to_idx[key1], sym_to_idx[key2]
                    result[zone_idx][text_idx] += input_data[key1][key2]
    if data_type == "infection":
        result = np.zeros(n_zone)
        for key in input_data:
            zone_idx = zoneid_to_idx[key]
            result[zone_idx] += input_data[key]
    return result

#function 1.24
#change the data to the matrix format
def change_to_matrix(data, zoneid_to_idx, sym_to_idx):
    data_result = list()
    for i in range(len(data)):
        combine1, combine2 = list(), list()
        combine3 = list()                                    #NEW
        mobility_text = data[i][0]
        x_infection_all = data[i][2]          #the x_days infection data
        day_order =  data[i][3] #NEW          the order of the day
        for j in range(round(len(mobility_text)*1.0/2)):
            mobility, text = mobility_text[2*j], mobility_text[2*j+1]
            x_infection =  x_infection_all[j]   #NEW
            new_mobility = to_matrix(zoneid_to_idx, sym_to_idx, mobility, "mobility")
            new_text = to_matrix(zoneid_to_idx, sym_to_idx, text, "text")
            combine1.append(new_mobility)
            combine1.append(new_text) 
            new_x_infection = to_matrix(zoneid_to_idx, sym_to_idx, x_infection, "infection") #NEW
            combine3.append(new_x_infection)   #NEW
        for j in range(len(data[i][1])):
            infection = data[i][1][j]                                                          
            new_infection = to_matrix(zoneid_to_idx, sym_to_idx, infection, "infection")
            combine2.append(new_infection)                                               
        data_result.append([combine1,combine2,combine3,day_order])    #mobility/text; infection_y; infection_x; day_order
    return data_result   

##################5.learn#####################
#function 1.25
def visual_loss(e_losses, vali_loss, test_loss):
    plt.figure(figsize=(4,3), dpi=300)
    x = range(len(e_losses))
    y1,y2,y3 = copy.copy(e_losses), copy.copy(vali_loss), copy.copy(test_loss)
    plt.plot(x,y1,linewidth=1, label="train")
    plt.plot(x,y2,linewidth=1, label="validate")
    plt.plot(x,y3,linewidth=1, label="test")
    plt.legend()
    plt.title('Loss decline on entire training/validation/testing data')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    #plt.savefig('final_f6.png',bbox_inches = 'tight')
    plt.show()

#function 1.26
def visual_loss_train(e_losses):
    plt.figure(figsize=(4,3), dpi=300)
    x = range(len(e_losses))
    y1 = copy.copy(e_losses)
    plt.plot(x,y1,linewidth=1, label="train")
    plt.legend()
    plt.title('Loss decline on entire training data')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    #plt.savefig('final_f6.png',bbox_inches = 'tight')
    plt.show()

Step 2: training epoch, training process, testing

#function 2.1
#normalize each column of the input mobility matrix as one
def normalize_column_one(input_matrix):
    column_sum = np.sum(input_matrix, axis=0)
    row_num, column_num = len(input_matrix), len(input_matrix[0])
    for i in range(row_num):
        for j in range(column_num):
             input_matrix[i][j] = input_matrix[i][j]*1.0/column_sum[j]
    return input_matrix

#function 2.2
#evalute the trained_model on validation or testing data.
def validate_test_process(trained_model, vali_test_data):
    criterion = nn.MSELoss()
    vali_test_y = [vali_test_data[i][1] for i in range(len(vali_test_data))]
    y_real = torch.tensor(vali_test_y)
    
    vali_test_x = [vali_test_data[i] for i in range(len(vali_test_data))]
    vali_test_x = convertAdj(vali_test_x)
    y_hat = trained_model.run_specGCN_lstm(vali_test_x)                                          
    loss = criterion(y_hat.float(), y_real.float())            ###Calculate the loss  
    return loss, y_hat, y_real 

#function 2.3
#convert the mobility matrix in x_batch in a following way
#normalize the flow between zones so that the in-flow of each zone is 1.
def convertAdj(x_batch):
    #x_batch：(n_batch, 0/1, 2*i+1)
    x_batch_new = copy.copy(x_batch)
    n_batch = len(x_batch)
    days = round(len(x_batch[0][0])/2)
    for i in range(n_batch):
        for j in range(days):
            mobility_matrix = x_batch[i][0][2*j]
            x_batch_new[i][0][2*j] = normalize_column_one(mobility_matrix)   #20210818
    return x_batch_new

#function 2.4
#a training epoch
def train_epoch_option(model, opt, criterion, trainX_c, trainY_c, batch_size):  
    model.train()
    losses = []
    batch_num = 0
    for beg_i in range(0, len(trainX_c), batch_size):
        batch_num += 1
        if batch_num % 16 ==0:
            print ("batch_num: ", batch_num, "total batch number: ", int(len(trainX_c)/batch_size))
        x_batch = trainX_c[beg_i:beg_i+batch_size]        
        y_batch = torch.tensor(trainY_c[beg_i:beg_i+batch_size])   
        opt.zero_grad()
        x_batch = convertAdj(x_batch)   #conduct the column normalization
        y_hat = model.run_specGCN_lstm(x_batch)                          ###Attention
        loss = criterion(y_hat.float(), y_batch.float()) #MSE loss
        #opt.zero_grad()
        loss.backward()
        opt.step()
        losses.append(loss.data.numpy())
    return sum(losses)/float(len(losses)), model

#function 2.5
#multiple training epoch
def train_process(train_data, lr, num_epochs, net, criterion, bs, vali_data, test_data):
    opt = optim.Adam(net.parameters(), lr, betas = (0.9,0.999), weight_decay=0) 
    train_y = [train_data[i][1] for i in range(len(train_data))]
    e_losses = list()
    e_losses_vali = list()
    e_losses_test = list()
    time00 = time.time()
    for e in range(num_epochs):
        time1 = time.time()
        print ("current epoch: ",e, "total epoch: ", num_epochs)
        number_list = list(range(len(train_data)))       
        random.shuffle(number_list)
        trainX_sample = [train_data[number_list[j]] for j in range(len(number_list))]
        trainY_sample = [train_y[number_list[j]] for j in range(len(number_list))]
        loss, net =  train_epoch_option(net, opt, criterion, trainX_sample, trainY_sample, bs)  
        print ("train loss", loss*infection_normalize_ratio*infection_normalize_ratio)
        e_losses.append(loss*infection_normalize_ratio*infection_normalize_ratio)
        
        loss_vali, y_hat_vali, y_real_vali = validate_test_process(net, vali_data) 
        loss_test, y_hat_test, y_real_test = validate_test_process(net, test_data)
        e_losses_vali.append(float(loss_vali)*infection_normalize_ratio*infection_normalize_ratio)
        e_losses_test.append(float(loss_test)*infection_normalize_ratio*infection_normalize_ratio)
        
        print ("validate loss", float(loss_vali)*infection_normalize_ratio*infection_normalize_ratio)
        print ("test loss", float(loss_test)*infection_normalize_ratio*infection_normalize_ratio)
        if e>=2 and (e+1)%10 ==0:
            visual_loss(e_losses, e_losses_vali, e_losses_test)     
            visual_loss_train(e_losses) 
        time2 = time.time()
        print ("running time for this epoch:", time2 - time1)
        time01 = time.time()
        print ("---------------------------------------------------------------")
        print ("---------------------------------------------------------------")
        #print ("total running time until now:", time01 - time00)
        #print ("------------------------------------------------")
        #print("specGCN_weight", net.specGCN.layer1.W)
        #print("specGCN_weight_grad", net.specGCN.layer1.W.grad)
        #print ("------------------------------------------------")
        #print("memory decay matrix", net.v)
        #print("memory decay matrix grad", net.v.grad)
        #print ("------------------------------------------------")
        #print ("lstm weight", net.lstm.all_weights[0][0])
        #print ("lstm weight grad", net.lstm.all_weights[0][0].grad)
        #print ("------------------------------------------------")
        #print ("fc1.weight", net.fc1.weight)
        #print ("fc1 weight grd", net.fc1.weight.grad)
        #print ("---------------------------------------------------------------")
        #print ("---------------------------------------------------------------")
    return e_losses, net

Step 3: models

#function 3.1
def read_data():
    jcode23 = list(read_tokyo_23()["JCODE"])                    #1.1 get the tokyo 23 zone shapefile
    all_mobility = read_mobility_data(jcode23)                  #1.2 read the mobility data
    all_text = read_text_data(jcode23)                          #1.3 read the text data
    all_infection, all_infection_cum = read_infection_data(jcode23)                #1.4 read the infection data
    
    #smooth the data using 7-days average
    window_size = WINDOW_SIZE                 #20210818
    dateList = generate_dateList()  #20210818
    all_mobility = mob_inf_smooth(all_mobility, window_size, dateList) #20210818
    all_infection = mob_inf_smooth(all_infection, window_size, dateList)  #20210818
    
    #smooth, user, min-max.
    point_json = read_point_json()                           #20210821
    all_text = normalize_text_user(all_text, point_json)       #20210821
    all_text = text_smooth(all_text, window_size, dateList) #20210818
    all_text = min_max_text_data(all_text,jcode23)                 #20210820
    
    x_days, y_days =  X_day, Y_day
    mobility_date_cut = sort_date_2(all_mobility, START_DATE, x_days, END_DATE, y_days)
    all_x_y = ensemble(all_mobility, all_text, all_infection, x_days, y_days, mobility_date_cut)
    train_original, validate_original, test_original, train_list, validation_list =\
    split_data_3(all_x_y,train_ratio,validate_ratio)  
    zone_dict, text_dict = get_zone_text_to_idx(all_infection)                       #get zone_idx, text_idx
    train_x_y = change_to_matrix(train_original, zone_dict, text_dict)                   #get train
    print ("train_x_y_shape",len(train_x_y),"train_x_y_shape[0]",len(train_x_y[0]))
    validate_x_y = change_to_matrix(validate_original, zone_dict, text_dict)             #get validate
    test_x_y = change_to_matrix(test_original, zone_dict, text_dict)                     #get test
    
    print (len(train_x_y))  #300
    print (len(train_x_y[0][0])) #14
    print (np.shape(train_x_y[0][0][0])) #(23,23)
    print (np.shape(train_x_y[0][0][1])) #(23,43)
    #print ("---------------------------------finish data reading and preprocessing------------------------------------")
    return train_x_y, validate_x_y, test_x_y, all_mobility, all_infection, train_original, validate_original, test_original, train_list, validation_list

#function 3.2
#train the model
def model_train(train_x_y, vali_data, test_data):
    #3.2.1 define the model
    input_dim_1, hidden_dim_1, out_dim_1, hidden_dim_2 = len(train_x_y[0][0][1][1]),\
    HIDDEN_DIM_1, OUT_DIM_1, HIDDEN_DIM_2 
    dropout_1, alpha_1, N = DROPOUT, ALPHA, len(train_x_y[0][0][1])
    G_L_Model = SpecGCN_LSTM(X_day, Y_day, input_dim_1, hidden_dim_1, out_dim_1, hidden_dim_2, dropout_1,N)         ###Attention
    #3.2.2 train the model
    num_epochs, batch_size, learning_rate = NUM_EPOCHS, BATCH_SIZE, LEARNING_RATE                                                 #model train
    criterion = nn.MSELoss() 
    e_losses, trained_model = train_process(train_x_y, learning_rate, num_epochs, G_L_Model, criterion, batch_size,\
                          vali_data, test_data)
    return e_losses, trained_model

#function 3.3
#evaluate the error on validation (or testing) data.
def validate_test_process(trained_model, vali_test_data):
    criterion = nn.MSELoss()
    vali_test_y = [vali_test_data[i][1] for i in range(len(vali_test_data))]
    y_real = torch.tensor(vali_test_y)
    vali_test_x = [vali_test_data[i] for i in range(len(vali_test_data))]
    vali_test_x = convertAdj(vali_test_x)
    y_hat = trained_model.run_specGCN_lstm(vali_test_x)                                  ###Attention              
    loss = criterion(y_hat.float(), y_real.float())
    return loss, y_hat, y_real 

Step 4. Model implementation

#4.1
#read the data
train_x_y, validate_x_y, test_x_y, all_mobility, all_infection, \
train_original, validate_original, test_original, train_list, validation_list =\
read_data()
#train_x_y, validate_x_y, test_x_y = normalize(train_x_y, validate_x_y, test_x_y)

#train_x_y = train_x_y[0:30]
print (len(train_x_y))
print ("---------------------------------finish data preparation------------------------------------")

train_x_y_shape 210 train_x_y_shape[0] 4
210
42
(23, 23)
(23, 8)
210
---------------------------------finish data preparation------------------------------------

#4.2
#train the model
e_losses, trained_model = model_train(train_x_y, validate_x_y, test_x_y)
print ("---------------------------finish model training-------------------------")

current epoch:  0 total epoch:  100
batch_num:  16 total batch number:  26
train loss 476.3313110365912
validate loss 316.41948968172073
test loss 510.28795540332794
running time for this epoch: 10.141513347625732
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  1 total epoch:  100
batch_num:  16 total batch number:  26
train loss 443.7961191352871
validate loss 312.63310462236404
test loss 504.14711236953735
running time for this epoch: 8.608564615249634
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  2 total epoch:  100
batch_num:  16 total batch number:  26
train loss 434.52891728116407
validate loss 290.0681085884571
test loss 460.68359166383743
running time for this epoch: 8.666988611221313
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  3 total epoch:  100
batch_num:  16 total batch number:  26
train loss 393.3782106334412
validate loss 219.63214501738548
test loss 374.00949746370316
running time for this epoch: 8.833451986312866
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  4 total epoch:  100
batch_num:  16 total batch number:  26
train loss 332.2398852712164
validate loss 156.07528388500214
test loss 303.96489426493645
running time for this epoch: 8.79914379119873
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  5 total epoch:  100
batch_num:  16 total batch number:  26
train loss 287.9430645973318
validate loss 116.11597612500191
test loss 250.86157023906708
running time for this epoch: 8.834282875061035
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  6 total epoch:  100
batch_num:  16 total batch number:  26
train loss 258.9216363547301
validate loss 95.5030508339405
test loss 214.10563960671425
running time for this epoch: 8.923089742660522
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  7 total epoch:  100
batch_num:  16 total batch number:  26
train loss 233.06636010400123
validate loss 84.08457972109318
test loss 182.92322754859924
running time for this epoch: 8.91145944595337
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  8 total epoch:  100
batch_num:  16 total batch number:  26
train loss 234.86291860333748
validate loss 80.60595951974392
test loss 160.19336879253387
running time for this epoch: 8.9170560836792
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  9 total epoch:  100
batch_num:  16 total batch number:  26
train loss 200.29340806003248
validate loss 82.17720314860344
test loss 139.65962454676628

running time for this epoch: 9.251115798950195
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  10 total epoch:  100
batch_num:  16 total batch number:  26
train loss 187.8576476506337
validate loss 87.5534676015377
test loss 123.59650805592537
running time for this epoch: 8.919320583343506
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  11 total epoch:  100
batch_num:  16 total batch number:  26
train loss 178.68668573943
validate loss 94.40804831683636
test loss 111.99604719877243
running time for this epoch: 8.962557315826416
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  12 total epoch:  100
batch_num:  16 total batch number:  26
train loss 175.3545231703255
validate loss 101.36696510016918
test loss 103.56266051530838
running time for this epoch: 8.999536037445068
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  13 total epoch:  100
batch_num:  16 total batch number:  26
train loss 187.3888414249652
validate loss 111.26830242574215
test loss 94.60089728236198
running time for this epoch: 9.01688528060913
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  14 total epoch:  100
batch_num:  16 total batch number:  26
train loss 179.9779536237043
validate loss 119.87223289906979
test loss 88.36002089083195
running time for this epoch: 9.071593761444092
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  15 total epoch:  100
batch_num:  16 total batch number:  26
train loss 170.13448876915157
validate loss 130.22586703300476
test loss 82.26456120610237
running time for this epoch: 9.085334539413452
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  16 total epoch:  100
batch_num:  16 total batch number:  26
train loss 178.59021745001274
validate loss 136.36822812259197
test loss 78.91177199780941
running time for this epoch: 8.9893639087677
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  17 total epoch:  100
batch_num:  16 total batch number:  26
train loss 159.62338424287736
validate loss 147.22063206136227
test loss 74.22317285090685
running time for this epoch: 9.053529739379883
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  18 total epoch:  100
batch_num:  16 total batch number:  26
train loss 158.16696988487686
validate loss 149.91508796811104
test loss 72.87256419658661
running time for this epoch: 9.112024545669556
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  19 total epoch:  100
batch_num:  16 total batch number:  26
train loss 156.7419578700706
validate loss 156.02458268404007
test loss 70.49043662846088

running time for this epoch: 9.318931341171265
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  20 total epoch:  100
batch_num:  16 total batch number:  26
train loss 155.77087420280333
validate loss 161.556638777256
test loss 68.50153673440218
running time for this epoch: 9.155923843383789
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  21 total epoch:  100
batch_num:  16 total batch number:  26
train loss 155.25259137912482
validate loss 165.85459932684898
test loss 66.94680079817772
running time for this epoch: 9.146969079971313
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  22 total epoch:  100
batch_num:  16 total batch number:  26
train loss 160.940309331097
validate loss 167.96352341771126
test loss 66.10800512135029
running time for this epoch: 9.09523868560791
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  23 total epoch:  100
batch_num:  16 total batch number:  26
train loss 150.76710014707513
validate loss 172.75340855121613
test loss 64.59100171923637
running time for this epoch: 9.159383296966553
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  24 total epoch:  100
batch_num:  16 total batch number:  26
train loss 151.6816983240898
validate loss 174.2273010313511
test loss 63.970061019062996
running time for this epoch: 9.17379379272461
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  25 total epoch:  100
batch_num:  16 total batch number:  26
train loss 150.75257875853114
validate loss 177.11104825139046
test loss 62.98517808318138
running time for this epoch: 9.101733446121216
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  26 total epoch:  100
batch_num:  16 total batch number:  26
train loss 152.79012309663273
validate loss 180.32681196928024
test loss 62.0611198246479
running time for this epoch: 9.168591260910034
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  27 total epoch:  100
batch_num:  16 total batch number:  26
train loss 163.49013403265013
validate loss 179.20246347784996
test loss 61.97573151439428
running time for this epoch: 9.232530355453491
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  28 total epoch:  100
batch_num:  16 total batch number:  26
train loss 147.97617517480697
validate loss 178.42361703515053
test loss 61.7482652887702
running time for this epoch: 9.11443305015564
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  29 total epoch:  100
batch_num:  16 total batch number:  26
train loss 150.48915537557116
validate loss 183.26513469219208
test loss 60.542039573192596

running time for this epoch: 9.433903455734253
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  30 total epoch:  100
batch_num:  16 total batch number:  26
train loss 162.94928379701793
validate loss 182.82709643244743
test loss 60.337185859680176
running time for this epoch: 9.224517822265625
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  31 total epoch:  100
batch_num:  16 total batch number:  26
train loss 159.94230647467904
validate loss 183.9262805879116
test loss 59.779928997159004
running time for this epoch: 9.090910196304321
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  32 total epoch:  100
batch_num:  16 total batch number:  26
train loss 147.18292815679754
validate loss 187.0831847190857
test loss 58.95869340747595
running time for this epoch: 9.15089464187622
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  33 total epoch:  100
batch_num:  16 total batch number:  26
train loss 148.72830671568713
validate loss 185.70689484477043
test loss 58.943708427250385
running time for this epoch: 9.22252082824707
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  34 total epoch:  100
batch_num:  16 total batch number:  26
train loss 147.76298396841244
validate loss 183.39501693844795
test loss 59.02715492993593
running time for this epoch: 9.108333826065063
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  35 total epoch:  100
batch_num:  16 total batch number:  26
train loss 150.39321759508718
validate loss 183.98793414235115
test loss 58.68158768862486
running time for this epoch: 9.148385047912598
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  36 total epoch:  100
batch_num:  16 total batch number:  26
train loss 148.26330761597663
validate loss 181.27908930182457
test loss 58.91520995646715
running time for this epoch: 9.207011699676514
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  37 total epoch:  100
batch_num:  16 total batch number:  26
train loss 145.02965914154495
validate loss 180.71265891194344
test loss 58.76920185983181
running time for this epoch: 9.133398056030273
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  38 total epoch:  100
batch_num:  16 total batch number:  26
train loss 143.91563875876642
validate loss 180.70640042424202
test loss 58.47834050655365
running time for this epoch: 9.193736553192139
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  39 total epoch:  100
batch_num:  16 total batch number:  26
train loss 144.21410548190275
validate loss 183.4746077656746
test loss 57.737152092158794

running time for this epoch: 9.524691820144653
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  40 total epoch:  100
batch_num:  16 total batch number:  26
train loss 141.980093213971
validate loss 179.88832667469978
test loss 58.07572975754738
running time for this epoch: 9.139913320541382
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  41 total epoch:  100
batch_num:  16 total batch number:  26
train loss 151.65945135608868
validate loss 177.65801399946213
test loss 58.29731002449989
running time for this epoch: 9.181942701339722
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  42 total epoch:  100
batch_num:  16 total batch number:  26
train loss 146.04028659286323
validate loss 182.53829330205917
test loss 57.07954056560993
running time for this epoch: 9.256416320800781
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  43 total epoch:  100
batch_num:  16 total batch number:  26
train loss 139.25380884083333
validate loss 176.83474346995354
test loss 57.73075856268406
running time for this epoch: 9.167900085449219
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  44 total epoch:  100
batch_num:  16 total batch number:  26
train loss 140.0167724824752
validate loss 174.58830028772354
test loss 57.89735820144415
running time for this epoch: 9.207130432128906
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  45 total epoch:  100
batch_num:  16 total batch number:  26
train loss 144.76103604667716
validate loss 175.78868195414543
test loss 57.38675594329834
running time for this epoch: 9.224085330963135
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  46 total epoch:  100
batch_num:  16 total batch number:  26
train loss 153.8791177207949
validate loss 171.72716557979584
test loss 57.838596403598785
running time for this epoch: 9.192791223526001
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  47 total epoch:  100
batch_num:  16 total batch number:  26
train loss 152.20130948970714
validate loss 175.77365040779114
test loss 56.82146642357111
running time for this epoch: 9.240610361099243
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  48 total epoch:  100
batch_num:  16 total batch number:  26
train loss 138.5921094118169
validate loss 177.7612417936325
test loss 56.18000403046608
running time for this epoch: 9.255104064941406
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  49 total epoch:  100
batch_num:  16 total batch number:  26
train loss 140.61671224664207
validate loss 173.66578802466393
test loss 56.64445925503969

running time for this epoch: 9.553390264511108
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  50 total epoch:  100
batch_num:  16 total batch number:  26
train loss 144.16762093013082
validate loss 172.0130443572998
test loss 56.691872887313366
running time for this epoch: 9.19110918045044
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  51 total epoch:  100
batch_num:  16 total batch number:  26
train loss 138.29135279930026
validate loss 171.47723585367203
test loss 56.51847459375858
running time for this epoch: 9.234419107437134
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  52 total epoch:  100
batch_num:  16 total batch number:  26
train loss 136.68199673432994
validate loss 169.772207736969
test loss 56.62074312567711
running time for this epoch: 9.144103050231934
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  53 total epoch:  100
batch_num:  16 total batch number:  26
train loss 137.02137089610375
validate loss 169.82011497020721
test loss 56.29703402519226
running time for this epoch: 9.20944356918335
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  54 total epoch:  100
batch_num:  16 total batch number:  26
train loss 135.99105899360168
validate loss 170.74448987841606
test loss 55.92425353825092
running time for this epoch: 9.250838041305542
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  55 total epoch:  100
batch_num:  16 total batch number:  26
train loss 138.7400145607966
validate loss 167.78312623500824
test loss 56.253508664667606
running time for this epoch: 9.177749872207642
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  56 total epoch:  100
batch_num:  16 total batch number:  26
train loss 136.5500631638699
validate loss 167.6006056368351
test loss 55.97358103841543
running time for this epoch: 9.210955619812012
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  57 total epoch:  100
batch_num:  16 total batch number:  26
train loss 136.5340798575845
validate loss 164.25712034106255
test loss 56.376513093709946
running time for this epoch: 9.224737644195557
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  58 total epoch:  100
batch_num:  16 total batch number:  26
train loss 137.06537627580542
validate loss 166.93135723471642
test loss 55.57445343583822
running time for this epoch: 9.18230676651001
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  59 total epoch:  100
batch_num:  16 total batch number:  26
train loss 135.04675102282178
validate loss 169.09347847104073
test loss 55.004251189529896

running time for this epoch: 9.480749607086182
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  60 total epoch:  100
batch_num:  16 total batch number:  26
train loss 137.15637037185607
validate loss 168.31794753670692
test loss 54.886648431420326
running time for this epoch: 9.23010516166687
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  61 total epoch:  100
batch_num:  16 total batch number:  26
train loss 145.13714627052346
validate loss 166.840311139822
test loss 54.98516373336315
running time for this epoch: 9.169922113418579
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  62 total epoch:  100
batch_num:  16 total batch number:  26
train loss 132.10954028181732
validate loss 168.03106293082237
test loss 54.48381882160902
running time for this epoch: 9.216758728027344
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  63 total epoch:  100
batch_num:  16 total batch number:  26
train loss 135.0160452744199
validate loss 167.11879521608353
test loss 54.387226700782776
running time for this epoch: 9.243929147720337
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  64 total epoch:  100
batch_num:  16 total batch number:  26
train loss 133.4970114166262
validate loss 166.40527173876762
test loss 54.27509546279907
running time for this epoch: 9.19033169746399
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  65 total epoch:  100
batch_num:  16 total batch number:  26
train loss 133.27329458358386
validate loss 164.2954908311367
test loss 54.40089385956526
running time for this epoch: 9.213933229446411
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  66 total epoch:  100
batch_num:  16 total batch number:  26
train loss 130.85630217877528
validate loss 164.30804505944252
test loss 54.176636040210724
running time for this epoch: 9.253240823745728
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  67 total epoch:  100
batch_num:  16 total batch number:  26
train loss 130.01966997291203
validate loss 163.88118267059326
test loss 53.9748789742589
running time for this epoch: 9.1518394947052
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  68 total epoch:  100
batch_num:  16 total batch number:  26
train loss 144.54511080489115
validate loss 162.4966599047184
test loss 54.01516333222389
running time for this epoch: 9.186725378036499
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  69 total epoch:  100
batch_num:  16 total batch number:  26
train loss 145.68529100398774
validate loss 164.92830589413643
test loss 53.38613409548998

running time for this epoch: 9.540785789489746
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  70 total epoch:  100
batch_num:  16 total batch number:  26
train loss 130.14180792702567
validate loss 165.39927572011948
test loss 53.05574741214514
running time for this epoch: 9.183481693267822
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  71 total epoch:  100
batch_num:  16 total batch number:  26
train loss 134.91596147004103
validate loss 163.5919325053692
test loss 53.12379449605942
running time for this epoch: 9.190329551696777
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  72 total epoch:  100
batch_num:  16 total batch number:  26
train loss 132.65775266551861
validate loss 157.80534595251083
test loss 53.915190510451794
running time for this epoch: 9.216739892959595
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  73 total epoch:  100
batch_num:  16 total batch number:  26
train loss 143.22675566223484
validate loss 153.2671321183443
test loss 54.43395581096411
running time for this epoch: 9.179774284362793
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  74 total epoch:  100
batch_num:  16 total batch number:  26
train loss 130.43477751866536
validate loss 156.83362260460854
test loss 53.50074265152216
running time for this epoch: 9.180930137634277
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  75 total epoch:  100
batch_num:  16 total batch number:  26
train loss 133.95856682176668
validate loss 156.65408223867416
test loss 53.41346375644207
running time for this epoch: 9.243659973144531
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  76 total epoch:  100
batch_num:  16 total batch number:  26
train loss 136.31266675441077
validate loss 155.67442402243614
test loss 53.31302992999554
running time for this epoch: 9.173897981643677
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  77 total epoch:  100
batch_num:  16 total batch number:  26
train loss 127.77715852415119
validate loss 156.88341110944748
test loss 52.84423008561134
running time for this epoch: 9.218085289001465
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  78 total epoch:  100
batch_num:  16 total batch number:  26
train loss 138.79219759945516
validate loss 157.52607956528664
test loss 52.57401615381241
running time for this epoch: 9.24601125717163
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  79 total epoch:  100
batch_num:  16 total batch number:  26
train loss 127.61182815733332
validate loss 157.61494636535645
test loss 52.332207560539246

running time for this epoch: 9.464875936508179
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  80 total epoch:  100
batch_num:  16 total batch number:  26
train loss 126.3317482166544
validate loss 155.70353716611862
test loss 52.44011525064707
running time for this epoch: 9.192445039749146
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  81 total epoch:  100
batch_num:  16 total batch number:  26
train loss 130.07538595788733
validate loss 154.31416220963
test loss 52.48954053968191
running time for this epoch: 9.247196912765503
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  82 total epoch:  100
batch_num:  16 total batch number:  26
train loss 141.6577372044601
validate loss 150.77127143740654
test loss 52.874283865094185
running time for this epoch: 9.163023710250854
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  83 total epoch:  100
batch_num:  16 total batch number:  26
train loss 127.60991532424535
validate loss 153.76364812254906
test loss 52.14868579059839
running time for this epoch: 9.193692922592163
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  84 total epoch:  100
batch_num:  16 total batch number:  26
train loss 129.93308271419394
validate loss 152.72258780896664
test loss 52.13135853409767
running time for this epoch: 9.238938570022583
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  85 total epoch:  100
batch_num:  16 total batch number:  26
train loss 136.3730742768557
validate loss 148.79758469760418
test loss 52.515119314193726
running time for this epoch: 9.188837051391602
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  86 total epoch:  100
batch_num:  16 total batch number:  26
train loss 126.8317767729362
validate loss 152.90792100131512
test loss 51.576560363173485
running time for this epoch: 9.228025197982788
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  87 total epoch:  100
batch_num:  16 total batch number:  26
train loss 124.90778919999246
validate loss 150.21909028291702
test loss 51.82734224945307
running time for this epoch: 9.264434337615967
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  88 total epoch:  100
batch_num:  16 total batch number:  26
train loss 124.04419740455019
validate loss 150.9397104382515
test loss 51.56096536666155
running time for this epoch: 9.205691576004028
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  89 total epoch:  100
batch_num:  16 total batch number:  26
train loss 126.31217233353742
validate loss 147.30090275406837
test loss 51.97012796998024

running time for this epoch: 9.521195888519287
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  90 total epoch:  100
batch_num:  16 total batch number:  26
train loss 124.43246134487842
validate loss 147.69732020795345
test loss 51.725972443819046
running time for this epoch: 9.256670475006104
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  91 total epoch:  100
batch_num:  16 total batch number:  26
train loss 125.89917263809454
validate loss 147.38032594323158
test loss 51.606022752821445
running time for this epoch: 9.199053287506104
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  92 total epoch:  100
batch_num:  16 total batch number:  26
train loss 124.15611400717385
validate loss 145.34587040543556
test loss 51.753632724285126
running time for this epoch: 9.218266487121582
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  93 total epoch:  100
batch_num:  16 total batch number:  26
train loss 123.19512406571043
validate loss 145.08619904518127
test loss 51.59646272659302
running time for this epoch: 9.232929468154907
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  94 total epoch:  100
batch_num:  16 total batch number:  26
train loss 124.36744808736775
validate loss 145.51382511854172
test loss 51.33600439876318
running time for this epoch: 9.156519651412964
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  95 total epoch:  100
batch_num:  16 total batch number:  26
train loss 124.80208697659825
validate loss 144.81733553111553
test loss 51.27324257045984
running time for this epoch: 9.189637184143066
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  96 total epoch:  100
batch_num:  16 total batch number:  26
train loss 134.86004986420826
validate loss 143.37420463562012
test loss 51.336344331502914
running time for this epoch: 9.275665760040283
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  97 total epoch:  100
batch_num:  16 total batch number:  26
train loss 124.37929795985971
validate loss 146.60649932920933
test loss 50.60520954430103
running time for this epoch: 9.210157871246338
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  98 total epoch:  100
batch_num:  16 total batch number:  26
train loss 122.91087885387242
validate loss 145.55469155311584
test loss 50.56751426309347
running time for this epoch: 9.189215660095215
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  99 total epoch:  100
batch_num:  16 total batch number:  26
train loss 124.950779998606
validate loss 143.18478293716908
test loss 50.82986783236265

running time for this epoch: 9.540159940719604
---------------------------------------------------------------
---------------------------------------------------------------
---------------------------finish model training-------------------------

#4.3 
print (len(train_x_y))
print (len(validate_x_y))
print (len(test_x_y))
#4.3.1 model validation
validation_result, validate_hat, validate_real = validate_test_process(trained_model, validate_x_y)
print ("---------------------------------finish model validation------------------------------------")
print (len(validate_hat))
print (len(validate_real))
#4.3.2 model testing
#4.4. model test
test_result, test_hat, test_real = validate_test_process(trained_model, test_x_y)
print ("---------------------------------finish model testing------------------------------------")
print (len(test_real))
print (len(test_hat))

210
30
60
---------------------------------finish model validation------------------------------------
30
30
---------------------------------finish model testing------------------------------------
60
60

Step 5: Evaluation

#5.1 RMSE, MAPE, MAE, RMSLE
def RMSELoss(yhat,y):
    return float(torch.sqrt(torch.mean((yhat-y)**2)))

def MAPELoss(yhat,y):
    return float(torch.mean(torch.div(torch.abs(yhat-y), y)))

def MAELoss(yhat,y):
    return float(torch.mean(torch.div(torch.abs(yhat-y), 1)))

def RMSLELoss(yhat,y):
    log_yhat = torch.log(yhat+1)
    log_y = torch.log(y+1)
    return float(torch.sqrt(torch.mean((log_yhat-log_y)**2)))    

#compute RMSE
rmse_validate = list()
rmse_test = list()
for i in range(len(validate_x_y)):
    rmse_validate.append(float(RMSELoss(validate_hat[i],validate_real[i])))
for i in range(len(test_x_y)):
    rmse_test.append(float(RMSELoss(test_hat[i],test_real[i])))
print ("rmse_validate mean", np.mean(rmse_validate))
print ("rmse_test mean", np.mean(rmse_test))

#compute MAE
mae_validate = list()
mae_test = list()
for i in range(len(validate_x_y)):
    mae_validate.append(float(MAELoss(validate_hat[i],validate_real[i])))
for i in range(len(test_x_y)):
    mae_test.append(float(MAELoss(test_hat[i],test_real[i])))
    
print ("mae_validate mean", np.mean(mae_validate))
print ("mae_test mean", np.mean(mae_test))

#show RMSE and MAE together
mae_validate, rmse_validate, mae_test, rmse_test =\
np.array(mae_validate)*infection_normalize_ratio, np.array(rmse_validate)*infection_normalize_ratio,\
np.array(mae_test)*infection_normalize_ratio, np.array(rmse_test)*infection_normalize_ratio
print ("-----------------------------------------")
print ("mae_validate mean", round(np.mean(mae_validate),3), "     rmse_validate mean", round(np.mean(rmse_validate),3))
print ("mae_test mean", round(np.mean(mae_test),3), "         rmse_test mean", round(np.mean(rmse_test),3))
print ("-----------------------------------------")

rmse_validate mean 0.10595423569365642
rmse_test mean 0.0677339696149905
mae_validate mean 0.08706403018420329
mae_test mean 0.05028699764075757
-----------------------------------------
mae_validate mean 8.706      rmse_validate mean 10.595
mae_test mean 5.029          rmse_test mean 6.773
-----------------------------------------

print(validate_hat[0][Y_day-1])
print(torch.sum(validate_hat[0][Y_day-1]))
print(validate_real[0][Y_day-1])
print(torch.sum(validate_real[0][Y_day-1]))

tensor([0.1244, 0.2152, 0.4253, 0.5619, 0.2311, 0.3229, 0.2501, 0.4530, 0.3266,
        0.4036, 0.7544, 0.9895, 0.4030, 0.4669, 0.5601, 0.3026, 0.3381, 0.2646,
        0.5327, 0.6211, 0.6205, 0.3002, 0.5968], grad_fn=<SelectBackward>)
tensor(10.0647, grad_fn=<SumBackward0>)
tensor([0.0229, 0.0600, 0.0900, 0.2014, 0.0500, 0.0786, 0.1164, 0.1614, 0.0929,
        0.1157, 0.4357, 0.4164, 0.1000, 0.1350, 0.1729, 0.2386, 0.1150, 0.0964,
        0.2600, 0.2907, 0.2500, 0.3986, 0.3093], dtype=torch.float64)
tensor(4.2079, dtype=torch.float64)

x = range(len(rmse_validate))
plt.figure(figsize=(8,2),dpi=300)
l1 = plt.plot(x, np.array(rmse_validate), 'ro-',linewidth=0.8, markersize=1.2, label='RMSE')
l2 = plt.plot(x, np.array(mae_validate), 'go-',linewidth=0.8, markersize=1.2, label='MAE')
plt.xlabel('Date from the first day of validation',fontsize=12)
plt.ylabel("RMSE/MAE daily new cases",fontsize=10)
my_y_ticks = np.arange(0,2100, 500)
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)
plt.legend()
plt.grid()
plt.show()

x = range(len(mae_test))
plt.figure(figsize=(8,2),dpi=300)
l1 = plt.plot(x, np.array(rmse_test), 'ro-',linewidth=0.8, markersize=1.2, label='RMSE')
l2 = plt.plot(x, np.array(mae_test), 'go-',linewidth=0.5, markersize=1.2, label='MAE')
plt.xlabel('Date from the first day of test',fontsize=12)
plt.ylabel("RMSE/MAE Daily new cases",fontsize=10)
my_y_ticks = np.arange(0,2100, 500)
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)
plt.legend()
plt.grid()
plt.show()

Correlation

from scipy import stats
#validate
y_days = Y_day
validate_hat_sum = [float(torch.sum(validate_hat[i][y_days-1])) for i in range(len(validate_hat))]
validate_real_sum = [float(torch.sum(validate_real[i][y_days-1])) for i in range(len(validate_real))]
print ("the correlation between validation: ", stats.pearsonr(validate_hat_sum, validate_real_sum)[0])
#test
test_hat_sum = [float(torch.sum(test_hat[i][y_days-1])) for i in range(len(test_hat))]
test_real_sum = [float(torch.sum(test_real[i][y_days-1])) for i in range(len(test_real))]
print ("the correlation between test: ", stats.pearsonr(test_hat_sum, test_real_sum)[0])
#train
train_result, train_hat, train_real = validate_test_process(trained_model, train_x_y)
train_hat_sum = [float(torch.sum(train_hat[i][0])) for i in range(len(train_hat))]
train_real_sum = [float(torch.sum(train_real[i][0])) for i in range(len(train_real))]
print ("the correlation between train: ", stats.pearsonr(train_hat_sum, train_real_sum)[0])

the correlation between validation:  0.6333346377035749
the correlation between test:  -0.13212861027651843
the correlation between train:  0.9338105865897467

step 6. Visualization

y1List = [np.sum(list(train_original[i+1][1][Y_day-1].values())) for i in range(len(train_original)-1)]
y2List = [np.sum(list(validate_original[i][1][Y_day-1].values())) for i in range(len(validate_original))]
y2List_hat = [float(torch.sum(validate_hat[i][Y_day-1])) for i in range(len(validate_hat))]
y3List = [np.sum(list(test_original[i][1][Y_day-1].values())) for i in range(len(test_original))]
y3List_hat = [float(torch.sum(test_hat[i][Y_day-1])) for i in range(len(test_hat))]

#x1 = np.array(range(len(y1List)))
#x2 = np.array([len(y1List)+j for j in range(len(y2List))])
x1 = train_list
x2 = validation_list
x3 = np.array([len(y1List)+len(y2List)+j for j in range(len(y3List))])

plt.figure(figsize=(8,2),dpi=300)
l1 = plt.plot(x1[0: len(y1List)], np.array(y1List)*infection_normalize_ratio, 'ro-',linewidth=0.8, markersize=2.0, label='train')
l2 = plt.plot(x2, np.array(y2List)*infection_normalize_ratio, 'go-',linewidth=0.8, markersize=2.0, label='validate')
l3 = plt.plot(x2, np.array(y2List_hat)*infection_normalize_ratio, 'g-',linewidth=2, markersize=0.1, label='validate_predict')
l4 = plt.plot(x3, np.array(y3List)*infection_normalize_ratio, 'bo-',linewidth=0.8, markersize=2, label='test')
l5 = plt.plot(x3, np.array(y3List_hat)*infection_normalize_ratio, 'b-',linewidth=2, markersize=0.1, label='test_predict')
#plt.xlabel('Date from the first day of 2020/4/1',fontsize=12)
plt.ylabel("Daily infection cases",fontsize=10)
my_y_ticks = np.arange(0,2100, 500)
my_x_ticks = list()
summary = 0 
my_x_ticks.append(summary) 
for i in range(5):
    summary += 60
    my_x_ticks.append(summary) 
plt.xticks(my_x_ticks)
plt.yticks(my_y_ticks) 
plt.xticks(fontsize=8)
plt.yticks(fontsize=12)
plt.title("SpectralGCN")
plt.legend()
plt.grid()
#plt.savefig('sg_peak4_21_21_1feature_0005.pdf',bbox_inches = 'tight')
plt.show()

step 7: Regional visualization

def getPredictionPlot(k):
    #location k
    x_k = [i for i in range(len(test_real))]
    real_k = [test_real[i][y_days-1][k] for i in range(len(test_real))]
    predict_k = [test_hat[i][y_days-1][k] for i in range(len(test_hat))]
    plt.figure(figsize=(4,2.5), dpi=300)
    l1 = plt.plot(x_k, np.array(real_k)*infection_normalize_ratio, 'ro-',linewidth=0.8, markersize=2.0, label='real',alpha = 0.8)
    l2 = plt.plot(x_k, np.array(predict_k)*infection_normalize_ratio, 'o-',color='black',linewidth=0.8, markersize=2.0, alpha = 0.8, label='predict')
    #plt.xlabel('Date from the first day of 2020/4/1',fontsize=12)
    #plt.ylabel("Daily infection cases",fontsize=10)
    my_y_ticks = np.arange(0,100,40)
    my_x_ticks = list()
    summary = 0 
    my_x_ticks.append(summary) 
    for i in range(6):
        summary += 10
        my_x_ticks.append(summary) 
    plt.xticks(my_x_ticks)
    plt.yticks(my_y_ticks) 
    plt.xticks(fontsize = 14)
    plt.yticks(fontsize = 14)
    plt.title("Real and predict daily infection for region "+str(k))
    plt.legend()
    plt.grid()
    #plt.savefig('sg_peak4_21_21_1feature_0005.pdf',bbox_inches = 'tight')
    plt.show()

for i in range(23):
    getPredictionPlot(i)